package hw2.memdir;

import java.io.*;
import java.security.*;
import java.security.spec.*;
import javax.crypto.*;
import javax.crypto.spec.*;

/**
 * 
 * Static method getAESkey used to create an AES key from a provided password
 * main used to create a random salt combined with the password to generate an
 * AES key
 * 
 */

/**
 * references: some guidance provided by
 * www.owasp.org/index.php/Using_the_Java_Cryptographic_Extensions#SecureRandom
 * stackoverflow.com/questions/992019/java-256bit-aes-encryption
 */
public class KeyGen
{

    /**
     * Main used only to generate salt.bin and proof-of-concept of AES key
     * generation
     * 
     * @param args
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws UnsupportedEncodingException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     * @throws InvalidParameterSpecException
     */
    public static void main( String[] args ) throws NoSuchAlgorithmException,
            InvalidKeySpecException, NoSuchPaddingException,
            InvalidKeyException, IllegalBlockSizeException,
            BadPaddingException, UnsupportedEncodingException,
            InvalidParameterSpecException {
       // SecretKeyFactory factory =
        // SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
        // SecureRandom secRand = SecureRandom.getInstance( "SHA1PRNG");
        // char[] password = "kruSEZweRULES".toCharArray();
        //
        // byte[] salt = readFile("salt.bin");
        // if(salt == null) {
        // salt = new byte[8];
        // secRand.nextBytes( salt );
        // writeFile("salt.bin", salt);
        // }
        // System.out.println("salt: " + new String(salt));

        // alternate way of generating secret
        // byte[] key = (new String(salt) + new
        // String(password)).getBytes("UTF-8");
        // MessageDigest sha = MessageDigest.getInstance("SHA-1");
        // key = sha.digest(key);
        // key = Arrays.copyOf(key, 16); // use only first 128 bit
        //
        // SecretKeySpec secret = new SecretKeySpec(key, "AES");

        // KeySpec spec = new PBEKeySpec(password , salt, 1024, 128);
        // SecretKey tmp = factory.generateSecret(spec);
        // SecretKey secret = new SecretKeySpec(tmp.getEncoded(), "AES");
        // System.out.println("Secret: " + secret.getEncoded());

        // use iv previously created
        byte[] iv = readFile( "ivpart2.bin" );

        // generate an AES key with a pre-determined password
        SecretKey secret = KeyGen.getAESkey( "kruSEZweRULES" );

        // start encryption
        Cipher cipher = Cipher.getInstance( "AES/CBC/PKCS5Padding" );
        try {
            // if an iv already existed
            if( iv != null ) {
                cipher.init( Cipher.ENCRYPT_MODE, secret, new IvParameterSpec(
                        iv ) );
            } else { // use a new iv and store it in the file system
                cipher.init( Cipher.ENCRYPT_MODE, secret );
                iv = cipher.getIV();
                writeFile( "ivpart2.bin", iv );
            }
        } catch( InvalidAlgorithmParameterException e ) {
            e.printStackTrace();
        }

        // store a predetermined message
        byte[] ciphertext = cipher
                .doFinal( "Hello, World!".getBytes( "UTF-8" ) );
        System.out.println( "IV: " + new String( iv ) );
        System.out.println( "Ciphertext: " + new String( ciphertext ) );

        // generate a separate (but same) AES key to verify decryption works
        // when using same password
        SecretKey decSecret = KeyGen.getAESkey( "kruSEZweRULES" );
        Cipher aesDec = Cipher.getInstance( "AES/CBC/PKCS5Padding" );
        try {
            aesDec.init( Cipher.DECRYPT_MODE, decSecret, new IvParameterSpec(
                    iv ) );
        } catch( InvalidAlgorithmParameterException e ) {
            e.printStackTrace();
        }

        byte[] text = aesDec.doFinal( ciphertext );
        System.out.println( "Decrypted text is: " + new String( text ) );

    }

    /**
     * Given a password string, generate an 128-bit AES key.
     * 
     * @param password
     *            the password desired for encryption/decryption
     * @return the AES key
     */
    public static SecretKey getAESkey( String password ) {
        SecretKeyFactory factory = null;
        try {
            factory = SecretKeyFactory.getInstance( "PBKDF2WithHmacSHA1" );
        } catch( NoSuchAlgorithmException e1 ) {
            e1.printStackTrace();
        }

        // convert password to char array
        char[] pass = password.toCharArray();

        // check if salt.bin exists, if not create a new one
        SecretKey secret = null;
        byte[] salt = readFile( "salt.bin" );
        if( salt == null ) {
            salt = createSalt( "salt.bin" );
        }

        // Password passed encryption, combined with the salt 128 bits long
        KeySpec spec = new PBEKeySpec( pass, salt, 1024, 128 );
        SecretKey tmp = null;
        try {
            tmp = factory.generateSecret( spec );
        } catch( InvalidKeySpecException e ) {
            e.printStackTrace();
        }
        // use PBE key to create the AES key
        secret = new SecretKeySpec( tmp.getEncoded(), "AES" );

        // return the AES key
        return secret;
    }

    public static byte[] createSalt( String string ) {
        // get a new secure random number in case salt doesn't exist
        SecureRandom secRand = null;
        try {
            secRand = SecureRandom.getInstance( "SHA1PRNG" );
        } catch( NoSuchAlgorithmException e ) {
            e.printStackTrace();
        }
        byte[] newSalt = new byte[8];
        secRand.nextBytes( newSalt );
        writeFile( "salt.bin", newSalt );
        return newSalt;
    }

    /**
     * 
     * @param filename
     *            a file to open
     * @return byte array of file contents
     */
    private static byte[] readFile( String filename ) {
        filename = "data/" + filename;
        byte[] data = null;
        try {
            File file = new File( filename );
            FileInputStream fin = new FileInputStream( filename );
            data = new byte[(int) file.length()];
            fin.read( data );
            fin.close();
        } catch( Exception e ) {
            System.out.println( "File " + filename + " not found" );
            // e.printStackTrace();
        }
        return data;
    }

    /**
     * 
     * @param filename
     *            a file to write to
     * @param data
     *            the bytes to store in the file
     */
    private static void writeFile( String filename, byte[] data ) {
        filename = "data/" + filename;
        try {
            FileOutputStream fout = new FileOutputStream( filename );
            fout.write( data );
            // ObjectOutputStream oos = new ObjectOutputStream(fout);

            // oos.writeObject(data);
            // oos.close();
            fout.close();
        } catch( Exception e ) {
            e.printStackTrace();
        }
    }

}
